# export https_proxy=http://proxy.329509.xyz:7890
# export NVI_NOTIFY_IGNORE_TASK=1
import os

os.environ["https_proxy"] = "http://proxy.329509.xyz:7890"

import multiprocessing

import shutil

import torch
import cv2
import numpy as np
from time import sleep as time_sleep

from ultralytics import YOLO

from nvitop import select_devices

import tqdm


def generate_params(dir_path: str, det_name: str):
    return (dir_path, det_name)


def get_torch_device() -> torch.device:
    device_list = None
    while device_list is None or len(device_list) == 0:
        device_list = select_devices(
            min_count=1, max_count=1, min_free_memory="4GiB", sort=True
        )

        if len(device_list) == 1:
            break

        time_sleep(2)

    return torch.device(device_list[0])


# 随机延时,单位秒s
def random_sleep(min_time: float = 0, max_time: float = 0.1):
    time_sleep(np.random.uniform(min_time, max_time))


def handle_dir(dir_path: str, det_name: str):

    random_sleep(5, 20)

    device: torch.device = get_torch_device()

    yolo = YOLO("mt-v11xbest.pt")
    yolo.to(device)

    video_name = os.path.basename(dir_path)
    parent_dir = os.path.dirname(dir_path)
    base_save_dir = os.path.join(os.path.dirname(parent_dir), "Det", det_name)
    os.makedirs(base_save_dir, exist_ok=True)

    frame_save_dir = os.path.join(base_save_dir, video_name)
    if os.path.exists(frame_save_dir):
        shutil.rmtree(frame_save_dir)
    os.makedirs(frame_save_dir, exist_ok=True)

    img_dir = os.path.join(dir_path, "img1")
    if not os.path.exists(img_dir):
        return

    img_list = os.listdir(img_dir)
    img_list = [
        img_name.replace(".jpg", "")
        for img_name in img_list
        if img_name.endswith(".jpg")
    ]
    img_list.sort()
    img_list = [img_name + ".jpg" for img_name in img_list]

    frame_index = 0
    frame_count = len(img_list)
    for image_name in tqdm.tqdm(img_list):
        image_path = os.path.join(img_dir, image_name)

        frame_index += 1
        print(f"Processing frame {frame_index}/{frame_count}")

        frame = cv2.imread(image_path)

        # Filter the detections (e.g., based on confidence threshold)
        confidence_threshold = 0
        dets = []

        yolo_result = yolo.predict(frame, conf=0.1)
        boxes = yolo_result[0].boxes
        for box in boxes:
            # 获取边界框坐标
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            # 获取置信度
            confidence = box.conf[0].tolist()
            # 获取类别索引
            class_index = box.cls[0].tolist()
            class_index = int(class_index)
            # 获取类别名称
            class_name = yolo.names[class_index]

            round_count = 6
            x1, y1, x2, y2 = (
                round(x1, round_count),
                round(y1, round_count),
                round(x2, round_count),
                round(y2, round_count),
            )

            confidence = round(confidence, round_count)

            if confidence >= confidence_threshold:
                dets.append([x1, y1, x2, y2, confidence, class_index])

                x1_i, y1_i, x2_i, y2_i = map(int, [x1, y1, x2, y2])

                # Plot Detection Info
                cv2.rectangle(frame, (x1_i, y1_i), (x2_i, y2_i), (0, 255, 0), 2)

                cv2.putText(
                    frame,
                    f"[{class_index}] {class_name} {confidence:.2f}",
                    (x1_i, min(y1_i - 10, 0)),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.9,
                    (0, 255, 0),
                    2,
                )

        # Convert detections to numpy array (N X (x, y, x, y, conf, cls))
        dets = np.array(dets)

        # Plot Frame Info
        cv2.putText(
            frame,
            f"Frame: {frame_index:03}/{frame_count:03}    Video Name: {video_name}",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 255, 0),
            2,
        )

        img_save_path = os.path.join(frame_save_dir, image_name)
        txt_save_path = os.path.join(frame_save_dir, image_name.replace(".jpg", ".txt"))

        cv2.imwrite(img_save_path, frame)

        txt_text_list = []
        for det in dets:
            x1, y1, x2, y2, confidence, class_index = det
            txt_text_list.append(f"{x1} {y1} {x2} {y2} {confidence} {int(class_index)}")
        with open(txt_save_path, "w") as f:
            f.write("\n".join(txt_text_list).strip())


if __name__ == "__main__":
    multiprocessing.set_start_method("spawn")

    base_dir_path = "/home/konghaomin/Datasets/MaritimeTrack_Full_Same/DanceTrack/val"

    base_dir_path = os.path.abspath(base_dir_path)

    ls_dir = os.listdir(base_dir_path)

    dir_list = []
    for dir in ls_dir:
        dir_path = os.path.join(base_dir_path, dir)
        if not os.path.isdir(dir_path):
            continue

        dir_list.append(dir_path)

    params_list = []

    for dir in dir_list:
        params_list.append(generate_params(dir_path=dir, det_name="v11x"))

    with multiprocessing.Pool(64) as pool:
        pool.starmap(handle_dir, params_list)

    print("done")
